`include "./parameters.v"

module sha3(
	input				clk,
	input 				rst_n,
	input 				enable,
	input 	[5:0]		model,
	input				ram_read_busy,
	input				ram_write_busy,

	input 	[63:0]		data_in,

	output	reg	[95:0]	data_out_a,
	output	reg [95:0]	data_out_b,

	input				read_en,
	output 				read_req,
	output	reg			valid_o,
	output	reg			done,

	input				done_rej
	);

	parameter	IDLE	=	4'D0,	READ 	=	4'D1,	READING =	4'D2,	PADDING	=	4'D3,	WORKING	=	4'D4;
	parameter	WAIT	=	4'D5,	WRITE	=	4'D6,	SQUEEZE	=	4'D7,	FINISH	=	4'D8,	BUFFER 	=	4'd9;
	parameter	BUFFER1	=	4'd10;



	parameter 	SHA3_512_R	= 11'd576,	SHAKE256_R	= 11'd1088,	SHAKE128_R	= 11'd1344,	SHA3_256_R	= 11'd1088;

	parameter 	SHA3_512_O	= 11'd512,	SHAKE256_O	= 11'd1088,	SHAKE128_O	= 11'd1344,	SHA3_256_O	= 11'd256;

	reg [3:0]	state,nextstate;
	reg	last_r,last_w;
	reg readdone,writedone;
	reg reading_done;
	reg	[4:0]	read_cnt;//有效数据个数
	reg [4:0]	reading_cnt;//请求输入数据
	reg [4:0]	write_cnt;
	reg	[4:0]	read_round_cnt,write_round_cnt;
	reg	clear;

//////////////////////////////////////////////////////////////////////////////////////////////////////
	reg	[63:0]	hashin;
	reg hash_read_en;
	wire [1599:0]	data_buffer;
	reg	padding,absorb,squeeze;
	wire hashdone;

	keccak	keccak(	.clk 		(clk),
					.rst_n      (rst_n),
					.enable     (enable),
					.model      (model[5:4]),
					.data_in    (hashin),
					.data_buffer(data_buffer),
					.padding    (padding),
					.read_en 	(hash_read_en),
					.absorb     (absorb),
					.squeeze    (squeeze),
					.hashdone   (hashdone));

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			absorb <= 1'b0;
		else if(~enable)
			absorb <= 1'b0;
		else if(state == READ | state == PADDING | state == READING)
			absorb <= 1'b1;
		else
			absorb <= 1'b0;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			padding <= 1'b0;
		else if(~enable)
			padding <= 1'b1;
		else if(last_r)//8704整除，最后一轮不需要填充
			padding <= 1'b1;
		else
			padding <= 1'b0;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			hash_read_en <= 1'b0;
		else if(~enable)
			hash_read_en <= 1'b0;
		else if(read_en | state == PADDING)
			hash_read_en <= 1'b1;
		else
			hash_read_en <=1'b0;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			squeeze <= 1'b0;
		else if(~enable)
			squeeze <= 1'b0;
		else if(nextstate == SQUEEZE & state == WAIT)
			squeeze <= 1'b1;
		else
			squeeze <= 1'b0;
	end
//in_lenth////////////////////////////////////////////////////////////////////////////////////////////////////
	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			last_r <= 1'b0;
		else if(~enable)
			last_r <= 1'b0;
		else
		begin
			case(model)
				6'b110001:	last_r <= (read_round_cnt == 5'd4  & state == WAIT) ? 1'b1 : last_r;
				6'b110010:	last_r <= (read_round_cnt == 5'd7  & state == WAIT) ? 1'b1 : last_r;
				6'b110011:	last_r <= (read_round_cnt == 5'd10 & state == WAIT) ? 1'b1 : last_r;
				6'b110100:	last_r <= (read_round_cnt == 5'd4  & state == WAIT) ? 1'b1 : last_r;
				6'b110101:	last_r <= (read_round_cnt == 5'd7  & state == WAIT) ? 1'b1 : last_r;
				6'b110110:	last_r <= (read_round_cnt == 5'd10 & state == WAIT) ? 1'b1 : last_r;
				default:	last_r <= 1'b1;
			endcase // model
		end
	end

	always@(*)
	begin
		case(model)
			6'b110001:	readdone = last_r ? read_cnt == 5'd14	: read_cnt == 5'd17;
			6'b110010:	readdone = last_r ? read_cnt == 5'd11	: read_cnt == 5'd17;
			6'b110011:	readdone = last_r ? read_cnt == 5'd8	: read_cnt == 5'd17;
			6'b110100:	readdone = last_r ? read_cnt == 5'd11	: read_cnt == 5'd17;
			6'b110101:	readdone = last_r ? read_cnt == 5'd0	: read_cnt == 5'd17;
			6'b110110:	readdone = last_r ? read_cnt == 5'd9	: read_cnt == 5'd17;

			6'b000010:	readdone = read_cnt == 5'd7;
			6'b011110:	readdone = read_cnt == 5'd7;
			default :	readdone = read_cnt == 4'd3;
		endcase // model
	end

	always@(*)
	begin
		case(model)
			6'b110001:	reading_done = last_r ? reading_cnt == 5'd14	: reading_cnt == 5'd16;
			6'b110010:	reading_done = last_r ? reading_cnt == 5'd11	: reading_cnt == 5'd16;
			6'b110011:	reading_done = last_r ? reading_cnt == 5'd8		: reading_cnt == 5'd16;
			6'b110100:	reading_done = last_r ? reading_cnt == 5'd10	: reading_cnt == 5'd16;
			6'b110101:	reading_done = last_r ? reading_cnt == 5'd0		: reading_cnt == 5'd16;
			6'b110110:	reading_done = last_r ? reading_cnt == 5'd8		: reading_cnt == 5'd16;
			6'b000010:	reading_done = reading_cnt == 5'd7;
			6'b011110:	reading_done = reading_cnt == 5'd7;
			default :	reading_done = reading_cnt == 4'd3;
		endcase // model
	end

//hashin////////////////////////////////////////////////////////////////////////////////////////////////////
	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			hashin <= 64'b0;
		else if(read_en & (state == READ || state == READING))
			hashin <= data_in;	
		else if(state == PADDING & last_r)
		begin
			case(model[5:4])
				`SHA3256:	hashin <= {56'b0,8'h06};
				`SHA3512:	hashin <= {56'b0,8'h06};
				`SHAKE128:	hashin <= {40'b0,8'h1f,6'b0,model[3:2],6'b0,model[1:0]};//矩阵填充
				`SHAKE256:	begin
								if(model[3:0] == 4'b1111)
									hashin <= {48'b0,8'h1f,8'h8};
								else if(model[3:0] == 4'b1110)
									hashin <= {56'b0,8'h1f};
								else
									hashin <= {48'b0,8'h1f,5'b0,model[2:0]};
							end
			endcase // model[5:4]
		end
		else
			hashin <= 64'b0;
	end
//out_lenth////////////////////////////////////////////////////////////////////////////////////////////////////
		always@(*)
		begin
			case(model[5:4])
				`SHAKE128:	last_w <= 1'b0;
				`SHAKE256:	begin
								case(model[3:0])
									4'b1000,4'b1001,4'b1010,4'b1011:	last_w <= write_round_cnt == 1'b1;
									default:	last_w <= 1'b1;
								endcase // model[3:0]
							end
				default	:	last_w <= 1'b1;
			endcase // model
		end

	// assign re_writedone = ~writedone;

	always@(*)
	begin
		case(model[5:4])
			`SHAKE128:	writedone = write_cnt == 5'd28;
			`SHA3512:	writedone = write_cnt == 5'd4;
			`SHA3256:	writedone = write_cnt == 5'd2;
			`SHAKE256:	begin
							case(model[3:0])
								4'b1110:							writedone = write_cnt == 5'd2;
								4'b1000,4'b1001,4'b1010,4'b1011:	writedone = last_w ? write_cnt == 5'd5 : write_cnt == 5'd11;
								default:							writedone = write_cnt == 5'd16;
							endcase // model[3:0]
						end
		endcase // model
	end
//write////////////////////////////////////////////////////////////////////////////////////////////////////
	reg	[1343:0]	buffer;

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			buffer <= 1344'b0;
		else if(state == WAIT)
		begin
			case(model)
				6'b011000,6'b011001,6'b011010,6'b011011:	buffer <= last_w ? {data_buffer[1311:0],buffer[31:0]} : data_buffer[1343:0];//剩余32bit拼接（eta1=3）
				default:	buffer <= data_buffer[1343:0];
			endcase // model
		end
		else if((state == WRITE | state == SQUEEZE | state == BUFFER1) & ~writedone)
		begin
			case(model[5:4])
				`SHA3256,`SHA3512:	buffer <= {128'b0,buffer[1343:128]};
				`SHAKE128:			buffer <= {48'b0,buffer[1343:48]};
				`SHAKE256:	begin
								case(model[3:0])
									4'b1110:	buffer <= {128'b0,buffer[1343:128]};
									4'b1000,4'b1001,4'b1010,4'b1011:	buffer <= {96'b0,buffer[1343:96]};
									default:	buffer <= {64'b0,buffer[1343:64]};
								endcase // model[3:0]
							end
			endcase // model[3:0]endcase
		end	
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			begin	data_out_a <= 96'b0;	data_out_b <= 96'b0;	end
		else if((state == WRITE | state == SQUEEZE | state == BUFFER1) & ~writedone)
		begin
			case(model[5:4])
				`SHA3256,`SHA3512:	begin	data_out_a <= {32'b0,buffer[63:0]};	data_out_b <= {32'b0,buffer[127:64]};	end
				`SHAKE128:			begin	data_out_a <= {48'b0,buffer[48:0]};			data_out_b <= 96'b0;	end
				`SHAKE256:	begin
								case(model[3:0])
									4'b1110:							begin	data_out_a <= {32'b0,buffer[63:0]};	data_out_b <= {32'b0,buffer[127:64]};	end
									4'b1000,4'b1001,4'b1010,4'b1011:	begin	data_out_a <= buffer[95:0];			data_out_b <= 96'b0;	end
									default:							begin	data_out_a <= {32'b0,buffer[63:0]};	data_out_b <= 96'b0;	end
								endcase // model[3:0]
							end
			endcase // model
		end
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			valid_o <= 1'b0;
		else if((state == WRITE | state == SQUEEZE | state == BUFFER1) & ~writedone)
			valid_o <= 1'b1;
		else
			valid_o <= 1'b0;
	end
//////////////////////////////////////////////////////////////////////////////////////////////////////
	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			state <= IDLE;
		else
			state <= nextstate;
	end

	always@(*)
	begin
		if(~enable | done_rej)
			nextstate = IDLE;
		// else if(done_rej)
		// 	nextstate <= FINISH;
		else
		begin
			case(state)
				IDLE:	nextstate = BUFFER;
				BUFFER:	nextstate = ram_read_busy ? BUFFER : READING;
				READING:nextstate = reading_done ? READ : READING;
				READ:	nextstate = readdone ? PADDING : READ;
				PADDING:nextstate = WORKING;
				WORKING:nextstate = hashdone ? WAIT : WORKING;
				WAIT:	nextstate = last_r ? ( ram_write_busy ? WAIT : (last_w ? WRITE : SQUEEZE ) ) : READING ;
				WRITE:	nextstate = writedone ? FINISH : WRITE;
				SQUEEZE:nextstate = hashdone ? BUFFER1 : SQUEEZE;
				BUFFER1:nextstate = writedone ? WAIT : BUFFER1;
				FINISH:	nextstate = IDLE;
				default:nextstate = IDLE;
			endcase // state
		end
	end


	assign read_req = enable ? state == READING : 1'b0;

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			done <= 1'b0;
		else if(state == FINISH | done_rej)
			done <= 1'b1;
		else
			done <= 1'b0;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			clear <= 1'b0;
		else if(state == FINISH)
			clear <= 1'b1;
		else
			clear <= 1'b0;
	end
//////////////////////////////////////////////////////////////////////////////////////////////////////////////
	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			read_cnt <= 5'b0;
		else if(~enable)
			read_cnt <= 5'b0;
		else if(clear | state == WAIT)
			read_cnt <= 5'b0;
		else if(read_en & ~readdone)
			read_cnt <= read_cnt + 1'b1;
		else
			read_cnt <= read_cnt;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			reading_cnt <= 5'b0;
		else if(~enable)
			reading_cnt <= 5'b0;
		else if(clear | state == WAIT)
			reading_cnt <= 5'b0;
		else if(read_req)
			reading_cnt <= reading_cnt + 1'b1;
		else
			reading_cnt <= reading_cnt;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			write_cnt <= 5'b0;
		else if(clear || ~enable)
			write_cnt <= 5'b0;
		else if((state == WRITE | state == SQUEEZE | state == BUFFER1) & ~writedone)
			write_cnt <= write_cnt + 1'b1;
		else if(state == WAIT)
			write_cnt <= 5'b0;
		else	
			write_cnt <= write_cnt;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			read_round_cnt <= 4'b0;
		else if(clear || ~enable)
			read_round_cnt <= 4'b0;
		else if(state == WAIT)
			read_round_cnt <= read_round_cnt + 1'b1;
		else
			read_round_cnt <= read_round_cnt;
	end

	always@(posedge clk or negedge rst_n)
	begin
		if(!rst_n)
			write_round_cnt <= 4'b0;
		else if(clear || ~enable)
			write_round_cnt <= 4'b0;
		else if(state == BUFFER1 & nextstate == WAIT)
			write_round_cnt <= write_round_cnt + 1'b1;
		else
			write_round_cnt <= write_round_cnt;
	end




endmodule